# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: xguo
import time
import json
import google.protobuf.json_format as json_format

from coin.exchange.base.kr_rest.private_client_base import PrivateClientBase
from coin.exchange.okex_futures.kr_rest.futures_product import OkexFuturesProduct
from coin.exchange.okex_futures.kr_rest.currency import OkexCurrency
from coin.exchange.okex_futures_v3.kr_rest.private_client import (
    convert_okex_order_type_to_proto_side)
from coin.proto.coin_query_pb2 import (
    AccountBalance,
    AccountPosition,
    CurrencyBalance,
    ProductOrderBundle,
    ProductOrderElement,
    ProductPosition,
    ProductFillElement,
    ProductFillBundle,
)
from coin.exchange.okex_futures_swap.kr_rest.native_private_client import (
    OkexSwapNativePrivateClient, OkexFuturesOrderStatus)
from coin.exchange.okex_futures_v3.kr_rest.private_client import (
    convert_okex_exec_type_to_proto_fill_type)
from coin.exchange.okex_futures_swap.kr_rest.product_info import okex_futures_swap_product_info
from cc.appcoin2.support.order.py_parser import PyOrderParser
from coin2.exchange.order.system_pb2 import OrderSystemConfig
from coin2.exchange.symbology_pb2 import ProductConfig


class OkexFuturesSwapPrivateParser(object):
  @staticmethod
  def parse_balance(msg):
    """"
    {
      "info": [
        {
          "equity": "0.0000",
          "fixed_balance": "0.0000",
          "instrument_id": "BTC-USD-SWAP",
          "margin": "0.0000",
          "margin_frozen": "0.0000",
          "margin_mode": "",
          "margin_ratio": "0.0000",
          "realized_pnl": "0.0000",
          "timestamp": "2019-01-07T11:01:34.376Z",
          "total_avail_balance": "0.0000",
          "unrealized_pnl": "0.0000"
        },
        {
          "equity": "10.0000",
          "fixed_balance": "0.0000",
          "instrument_id": "EOS-USD-SWAP",
          "margin": "0.0000",
          "margin_frozen": "0.0000",
          "margin_mode": "crossed",
          "margin_ratio": "",
          "realized_pnl": "0.0000",
          "timestamp": "2019-01-07T06:16:51.000Z",
          "total_avail_balance": "10.0000",
          "unrealized_pnl": "0.0000"
        }
      }
    """
    balance_list = []
    for balance_info in msg['info']:
      product = OkexFuturesProduct.FromStrNativeProductV3(balance_info['instrument_id'])
      currency = product.base
      total = float(balance_info['equity'])
      hold = float(balance_info['margin'])
      available = total - hold
      balance = CurrencyBalance(
          currency=currency.currency,
          total=total,
          available=available,
          hold=hold,
      )
      balance_list.append(balance)
    return AccountBalance(
        exchange='Okex',
        each_balance=balance_list,
    )

  @staticmethod
  def parse_margin_balance(msg):
    currency_balance = {}
    for balance_info in msg['info']:
      currency = OkexCurrency.FromStrNativeCurrency(balance_info['currency'])
      total = float(balance_info['equity'])
      hold = float(balance_info['margin'])
      available = total - hold
      if currency.currency in currency_balance:
        balance = currency_balance[currency.currency]
        balance.total += total
        balance.hold += hold
        balance.available += available
      else:
        balance = CurrencyBalance(
            currency=currency.currency,
            total=total,
            available=available,
            hold=hold,
        )
        currency_balance[currency.currency] = balance
    balance_list = list(currency_balance.values())
    return AccountBalance(
        exchange='Okex',
        each_balance=balance_list,
    )

  @staticmethod
  def parse_position(pos_msg):
    """
    REST
    {
      "margin_mode": "crossed",
      "holding": [
        {
          "liquidation_price": "0.07",
          "position": "1",
          "avail_position": "1",
          "avg_cost": "210.24",
          "settlement_price": "210.34",
          "instrument_id": "BTC-USD-SWAP",
          "leverage": "10",
          "realized_pnl": "0.00307096",
          "side": "short",
          "timestamp": "2018-10-17T20:11:00.443Z"
        }
      ]
    }
    """
    assert pos_msg['margin_mode'] == 'crossed', pos_msg
    positions = {}
    for msg in pos_msg['holding']:
      product = OkexFuturesProduct.FromStrNativeProduct(msg['instrument_id'])
      if product in positions:
        product_position = positions[product]
      else:
        product_position = ProductPosition()
        product_position.symbol = product.symbol
        positions[product] = product_position

      if msg['side'] == 'short':
        product_position.short_position = -float(msg['position'])
        product_position.available_short_position = -float(msg['avail_position'])
        product_position.short_realized_pnl = float(msg['realized_pnl'])
        product_position.short_avg_entry_price = float(msg['avg_cost'])
      elif msg['side'] == 'long':
        product_position.long_position = float(msg['position'])
        product_position.available_long_position = float(msg['avail_position'])
        product_position.long_realized_pnl = float(msg['realized_pnl'])
        product_position.long_avg_entry_price = float(msg['avg_cost'])
      elif msg['side'] in ('1', '2'):
        # Happens when this product position is zero.
        assert int(msg['position']) == 0, msg
        product_position.short_position = 0.
        product_position.short_realized_pnl = 0.
        product_position.short_avg_entry_price = 0.
        product_position.long_position = 0.
        product_position.long_realized_pnl = 0.
        product_position.long_avg_entry_price = 0.
      else:
        raise ValueError('Unknown side! %s' % msg)

    for product_position in positions.values():
      product_position.net_position = \
          product_position.long_position + product_position.short_position

    return AccountPosition(
        exchange='Okex',
        market_type='Futures',
        each_position=positions.values(),
    )

  @staticmethod
  def parse_rest_position(pos_msg):
    holding_list = [holding for each_pos in pos_msg for holding in each_pos['holding']]
    msg = {'margin_mode': 'crossed', 'holding': holding_list}
    account_position = OkexFuturesSwapPrivateParser.parse_position(msg)
    return account_position

  @staticmethod
  def parse_ws_product_position(pos_msg):
    """
    WS
    {
      "table": "swap/position",
      "data": [
        {
          "holding": [
            {
              "avail_position": "1.000",
              "avg_cost": "48.0000",
              "leverage": "10",
              "liquidation_price": "52.5359",
              "margin": "0.018",
              "position": "1.000",
              "realized_pnl": "-0.001",
              "settlement_price": "48.0000",
              "side": "short",
              "timestamp": "2018-11-29T02:37:01.963Z"
            }
          ],
          "instrument_id": "LTC-USD-SWAP",
          "margin_mode": "crossed"
        }
      ]
    }
    """
    assert pos_msg['table'] == 'swap/position', pos_msg
    assert len(pos_msg['data']) == 1, pos_msg
    assert pos_msg['data'][0]['margin_mode'] == 'crossed', pos_msg
    product = OkexFuturesProduct.FromStrNativeProduct(pos_msg['data'][0]['instrument_id'])
    product_position = ProductPosition(symbol=product.symbol,
                                       native_symbol=product.native_symbol)
    holding = pos_msg['data'][0]['holding']
    assert len(holding) <= 2, holding
    for msg in holding:
      if msg['side'] == 'short':
        product_position.short_position = -float(msg['position'])
        product_position.available_short_position = -float(msg['avail_position'])
        product_position.short_realized_pnl = float(msg['realized_pnl'])
        product_position.short_avg_entry_price = float(msg['avg_cost'])
      elif msg['side'] == 'long':
        product_position.long_position = float(msg['position'])
        product_position.available_long_position = float(msg['avail_position'])
        product_position.long_realized_pnl = float(msg['realized_pnl'])
        product_position.long_avg_entry_price = float(msg['avg_cost'])
      else:
        raise ValueError('Unknown side!')

    product_position.net_position = product_position.long_position + product_position.short_position
    return product_position

  @staticmethod
  def parse_orders(msg, product):
    """
    {
      "order_info": [
        {
          "client_oid": "",
          "contract_val": "10",
          "fee": "0.000000",
          "filled_qty": "0.0000",
          "instrument_id": "EOS-USD-SWAP",
          "order_id": "6a-4-3c6579727-0",
          "price": "2.739",
          "price_avg": "0.000",
          "size": "1.0000",
          "status": "0",
          "timestamp": "2019-01-08T07:28:50.983Z",
          "type": "4"
        }
      ]
    }

    """
    order_list = []
    for order_info in msg['order_info']:
      native_symbol = order_info['instrument_id']
      if native_symbol != product.native_symbol:
        return

      order = ProductOrderElement(
          symbol=product.symbol,
          order_id=order_info['order_id'],
          price=float(order_info['price']),
          qty=float(order_info['size']),
          side=convert_okex_order_type_to_proto_side(order_info['type']),
      )
      order_list.append(order)
    return ProductOrderBundle(
        exchange='Okex',
        symbol=product.symbol,
        each_order=order_list,
    )

  @staticmethod
  def parse_fills(order_list_msg, fill_msgs, product):
    assert order_list_msg['order_info'], order_list_msg
    fill_list = []
    for index, order_info in enumerate(order_list_msg['order_info']):
      native_symbol = order_info['instrument_id']
      if native_symbol != product.native_symbol:
        return

      assert int(order_info['status']) == OkexFuturesOrderStatus.FILLED
      fill_msg = fill_msgs[index]
      assert len(fill_msg) >= 1
      fill_msg = fill_msg[0]
      prod = OkexFuturesProduct.FromStrNativeProduct(fill_msg['instrument_id'])
      assert prod == product
      fee = -float(fill_msg['fee'])
      fill_type = convert_okex_exec_type_to_proto_fill_type(fill_msg['exec_type'])
      fill = ProductFillElement(symbol=product.symbol,
                                order_id=order_info['order_id'],
                                price=float(fill_msg['price']),
                                qty=float(fill_msg['order_qty']),
                                side=convert_okex_order_type_to_proto_side(order_info['type']),
                                fee=fee,
                                fee_currency=product.base.symbol,
                                fill_type=fill_type,
                                contract_val=float(order_info['contract_val']))
      fill_list.append(fill)
    return ProductFillBundle(
        exchange='Okex',
        symbol=product.symbol,
        each_fill=fill_list,
    )


class OkexFuturesSwapPrivateClient(PrivateClientBase):
  ProductType = OkexFuturesProduct

  def __init__(self, key_file):
    super().__init__()
    self.nprivc = OkexSwapNativePrivateClient(key_file=key_file)
    pi_bundle = okex_futures_swap_product_info._product_infos
    assert pi_bundle.HasField('mea')
    oss_config = OrderSystemConfig(
        mea=pi_bundle.mea,
        products=ProductConfig(norms=[pi.symbol for pi in pi_bundle.product_infos]))
    self._parser = PyOrderParser.from_mea_str(pi_bundle.mea)
    self._parser.add_product(oss_config.SerializeToString())

  def query_account_balance_impl(self):
    update = self.nprivc.query_accounts()
    msg = self._parser.parse_account_balance(json.dumps(update.msg))
    account_balance = AccountBalance()
    json_format.Parse(msg, account_balance)
    update.msg = account_balance
    return update

  def query_account_position_impl(self):
    update = self.nprivc.query_positions()
    msg = self._parser.parse_account_position(json.dumps(update.msg))
    account_position = AccountPosition()
    json_format.Parse(msg, account_position)
    update.msg = account_position
    return update

  def query_fills_impl(self, product):
    update = self.nprivc.query_order_list(status=OkexFuturesOrderStatus.FILLED,
                                          instrument_id=product.native_symbol,
                                          limit=50)
    fill_updates = []
    for order in update.msg['order_info']:
      native_symbol = order['instrument_id']
      if native_symbol != product.native_symbol:
        continue
      order_id = str(order['order_id'])
      fill_update = self.nprivc.query_transaction_details(order_id=order_id,
                                                          instrument_id=product.native_symbol)
      time.sleep(0.2)  # rate limit 10 times per 2 seconds
      fill_updates.append(fill_update.msg)
    update.msg = OkexFuturesSwapPrivateParser.parse_fills(update.msg, fill_updates, product)
    return update

  def query_submit_impl(self, product, side, price, qty):
    raise NotImplementedError()

  def query_cancel_impl(self, product, order_id):
    raise NotImplementedError()

  def query_cancel_product_impl(self, product):
    raise NotImplementedError()

  def query_account_fills_impl(self):
    raise NotImplementedError()

  def query_account_list_orders_impl(self):
    raise NotImplementedError()

if __name__ == '__main__':
  pv = OkexFuturesSwapPrivateClient('../../coin_key/view/prestoinvt00_swap/OKEx/view_v3_key.json')
  print(pv.query_account_balance_impl().msg)
  print(pv.query_account_position_impl().msg)
